import numpy as np
import os

# train knn model
def modelKnn(data_test, data_train, k):
    preLabel = []
    for sample_test in data_test:
        dist = []
        for sample_train in data_train:
            distSample = np.sqrt(np.sum(np.square(sample_test[:-1] - sample_train[:-1])))
            dist.append([distSample, sample_train[-1]])
        dist.sort()
        iterm = [d[-1] for d in dist[:k]]
        label = max(iterm, key=iterm.count)
        preLabel.append(label)
    y_test = data_test[:, -1]
    acc = np.mean(y_test == preLabel)
    return acc

# image to vector
def img2vector(filename):
    vector = np.zeros(32 * 32)
    index = 0
    with open(filename, 'r', encoding='utf-8') as fr:
        for i in range(32):
            line = fr.readline()
            for j in range(32):
                vector[index] = int(line[j])
                index += 1
    return vector

# load datasets
def loadDataSet(dataDir):
    dataList = os.listdir(dataDir)
    x_data = []
    y_data = []
    for fileName in dataList:
        label = fileName.split('_')[0]
        x = img2vector(dataDir + fileName)
        x_data.append(x)
        y_data.append(int(label))
    return x_data, y_data

# main program
if __name__ == '__main__':
    trainDir = 'datasets_ML2/Digits_data/trainingDigits/'
    testDir = 'datasets_ML2/Digits_data/testDigits/'

    x_train, y_train = loadDataSet(trainDir)
    x_test, y_test = loadDataSet(testDir)

    data_train = np.c_[x_train, y_train]
    data_test = np.c_[x_test, y_test]

    acc = modelKnn(data_test, data_train, k=3)
    print(acc)

# sklearn lib
from sklearn.neighbors import KNeighborsClassifier
model_knn = KNeighborsClassifier()
model_knn.fit(x_train, y_train)

print(model_knn.score(x_test, y_test))

predict_y = model_knn.predict(x_test)

from sklearn.metrics import confusion_matrix, precision_score, recall_score, f1_score, accuracy_score
print("confusion_matrix: \n", confusion_matrix(y_test, predict_y))
print("P: ", precision_score(y_test, predict_y, average=None))
print("R: ", recall_score(y_test, predict_y, average=None))
print("f1: ", f1_score(y_test, predict_y, average=None))
print("acc: ", accuracy_score(y_test, predict_y))
